import math

import numpy as np


def anchor_sizes(n_layers=4):
    s = np.linspace(0.01, 0.2, n_layers + 1)
    sizes = []
    for i in range(len(s) - 1):
        size = [s[i], math.sqrt(s[i] * s[i + 1])]
        sizes.append(size)

    return sizes


def anchor_boxes(feature_shape, image_shape, index=0, n_layers=4,
                 aspect_ratios=(1, 2, 0.5)):
    sizes = anchor_sizes(n_layers)[index]

    n_boxes = len(aspect_ratios) + 1

    image_height, image_width, _ = image_shape

    feature_height, feature_width, _ = feature_shape

    norm_height = image_height * sizes[0]
    norm_width = image_width * sizes[0]

    width_height = []

    for ar in aspect_ratios:
        box_width = norm_width * np.sqrt(ar)
        box_height = norm_height / np.sqrt(ar)
        width_height.append((box_width, box_height))

    box_width = image_width * sizes[1]
    box_height = image_height * sizes[1]
    width_height.append((box_width, box_height))

    width_height = np.array(width_height)

    grid_width = image_width / feature_width
    grid_height = image_height / feature_height

    start = grid_width * 0.5
    end = (feature_width - 0.5) * grid_width

    cx = np.linspace(start, end, feature_width)

    start = grid_height * 0.5
    end = (feature_height - 0.5) * grid_height

    cy = np.linspace(start, end, feature_height)

    cx_grid, cy_grid = np.meshgrid(cx, cy)
    # 由cx和cy生成坐标网格，cx_grid以cx为行向量重复，cy_grid重复列向量cy

    cx_grid = np.expand_dims(cx_grid, -1)
    cy_grid = np.expand_dims(cy_grid, -1)

    boxes = np.zeros((feature_height, feature_width, n_boxes, 4))

    boxes[..., 0] = np.tile(cx_grid, (1, 1, n_boxes))
    boxes[..., 1] = np.tile(cy_grid, (1, 1, n_boxes))

    boxes[..., 2] = width_height[:, 0]
    boxes[..., 3] = width_height[:, 1]

    boxes = centroid2minmax(boxes)
    boxes = np.expand_dims(boxes, axis=0)

    return boxes


def centroid2minmax(boxes):
    minmax = np.copy(boxes).astype(np.float)
    minmax[..., 0] = boxes[..., 0] - (0.5 * boxes[..., 2])
    minmax[..., 1] = boxes[..., 0] + (0.5 * boxes[..., 2])
    minmax[..., 2] = boxes[..., 1] - (0.5 * boxes[..., 3])
    minmax[..., 3] = boxes[..., 1] + (0.5 * boxes[..., 3])
    return minmax


def minmax2centroid(boxes):
    centroid = np.copy(boxes).astype(np.float)
    centroid[..., 0] = 0.5 * (boxes[..., 1] - boxes[..., 0])
    centroid[..., 0] += boxes[..., 0]
    centroid[..., 1] = 0.5 * (boxes[..., 3] - boxes[..., 2])
    centroid[..., 1] += boxes[..., 2]
    centroid[..., 2] = boxes[..., 1] - boxes[..., 0]
    centroid[..., 3] = boxes[..., 3] - boxes[..., 2]
    return centroid


def intersection(boxes1, boxes2):
    """Compute intersection of batch of boxes1 and boxes2

    Arguments:
        boxes1 (tensor): Boxes coordinates in pixels
        boxes2 (tensor): Boxes coordinates in pixels

    Returns:
        intersection_areas (tensor): intersection of areas of
            boxes1 and boxes2
    """
    m = boxes1.shape[0]  # The number of boxes in `boxes1`
    n = boxes2.shape[0]  # The number of boxes in `boxes2`

    xmin = 0
    xmax = 1
    ymin = 2
    ymax = 3

    boxes1_min = np.expand_dims(boxes1[:, [xmin, ymin]], axis=1)
    boxes1_min = np.tile(boxes1_min, reps=(1, n, 1))
    boxes2_min = np.expand_dims(boxes2[:, [xmin, ymin]], axis=0)
    boxes2_min = np.tile(boxes2_min, reps=(m, 1, 1))
    min_xy = np.maximum(boxes1_min, boxes2_min)

    boxes1_max = np.expand_dims(boxes1[:, [xmax, ymax]], axis=1)
    boxes1_max = np.tile(boxes1_max, reps=(1, n, 1))
    boxes2_max = np.expand_dims(boxes2[:, [xmax, ymax]], axis=0)
    boxes2_max = np.tile(boxes2_max, reps=(m, 1, 1))
    max_xy = np.minimum(boxes1_max, boxes2_max)

    side_lengths = np.maximum(0, max_xy - min_xy)

    intersection_areas = side_lengths[:, :, 0] * side_lengths[:, :, 1]
    return intersection_areas


def union(boxes1, boxes2, intersection_areas):
    m = boxes1.shape[0]  # number of boxes in boxes1
    n = boxes2.shape[0]  # number of boxes in boxes2

    xmin = 0
    xmax = 1
    ymin = 2
    ymax = 3

    width = (boxes1[:, xmax] - boxes1[:, xmin])
    height = (boxes1[:, ymax] - boxes1[:, ymin])
    areas = width * height
    boxes1_areas = np.tile(np.expand_dims(areas, axis=1), reps=(1, n))
    width = (boxes2[:, xmax] - boxes2[:, xmin])
    height = (boxes2[:, ymax] - boxes2[:, ymin])
    areas = width * height
    boxes2_areas = np.tile(np.expand_dims(areas, axis=0), reps=(m, 1))

    union_areas = boxes1_areas + boxes2_areas - intersection_areas
    return union_areas


def iou(boxes1, boxes2):
    """Compute IoU of batch boxes1 and boxes2
    使用min，max格式，x_min,x_max,y_min,y_max

    Arguments:
        boxes1 (tensor): Boxes coordinates in pixels
        boxes2 (tensor): Boxes coordinates in pixels

    Returns:
        iou (tensor): intersectiin of union of areas of
            boxes1 and boxes2
    """
    intersection_areas = intersection(boxes1, boxes2)
    union_areas = union(boxes1, boxes2, intersection_areas)
    # union 为0的时候 inter 肯定为0，所以分母为1也没有影响
    union_areas[union_areas == 0] = 1
    return intersection_areas / union_areas


def get_gt_data(iou_value, n_classes=4, anchors=None, labels=None,
                threshold=0.6):
    max_iou_per_gt = np.argmax(iou_value, axis=0)
    if threshold < 1.0:
        iou_gt_thresh = np.argwhere(iou_value > threshold)

        if iou_gt_thresh.size > 0:
            extra_anchors = iou_gt_thresh[:, 0]
            extra_classes = iou_gt_thresh[:, 1]

            extra_labels = labels[extra_classes]

            indexes = [max_iou_per_gt, extra_anchors]
            max_iou_per_gt = np.concatenate(indexes, axis=0)
            labels = np.concatenate([labels, extra_labels], axis=0)
    gt_mask = np.zeros((iou_value.shape[0], 4))
    gt_mask[max_iou_per_gt] = 1.0

    gt_class = np.zeros((iou_value.shape[0], n_classes))

    gt_class[:, 0] = 1  # 默认都是背景，0为背景
    gt_class[max_iou_per_gt, 0] = 0
    max_iou_col = np.reshape(max_iou_per_gt, (max_iou_per_gt.shape[0], 1))
    label_col = np.reshape(labels[:, 4], (labels.shape[0], 1)).astype(int)

    row_col = np.append(max_iou_col, label_col, axis=1)

    gt_class[row_col[:, 0], row_col[:, 1]] = 1.0

    gt_offset = np.zeros((iou_value.shape[0], 4))

    anchors = minmax2centroid(anchors)
    labels = minmax2centroid(labels)

    offsets1 = labels[:, 0:2] - anchors[max_iou_per_gt, 0:2]
    offsets1 /= anchors[max_iou_per_gt, 2:4]
    offsets1 /= 0.1  # \sigma_x ,\sigma_y

    offsets2 = np.log(labels[:, 2:4] / anchors[max_iou_per_gt, 2:4])
    offsets2 /= 0.2  # \sigma_h ,\sigma_w

    offsets = np.concatenate([offsets1, offsets2], axis=-1)

    gt_offset[max_iou_per_gt] = offsets
    return gt_class, gt_offset, gt_mask
